# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations
from os import path as osp
from typing import Optional
from typing import Tuple
from typing import Union

import sympy as sp
from sympy.parsing import sympy_parser as sp_parser

import ppsci
from ppsci.equation.pde import base

import hydra
import numpy as np
from omegaconf import DictConfig

import ppsci
import paddle
from ppsci.autodiff import hessian
from ppsci.autodiff import jacobian
from ppsci.utils import logger


class CNC2D(base.PDE):
    def __init__(
            self,
            pr: Union[float, str],
            ra: Union[float, str],
            dim: int,
            detach_keys: Optional[Tuple[str, ...]] = None,
    ):
        super().__init__()
        self.detach_keys = detach_keys
        self.dim = dim
        # self.time = time

        x, y, z = self.create_symbols("x y z")
        invars = (x, y)
        # if time:
        #     invars = (t,) + invars
        if dim == 3:
            invars += (z,)

        if isinstance(pr, str):
            print("pr is str")
            pr = sp_parser.parse_expr(pr)
            if isinstance(pr, sp.Symbol):
                invars += (pr,)

        if isinstance(ra, str):
            print("ra is str")
            ra = sp_parser.parse_expr(ra)
            if isinstance(ra, sp.Symbol):
                invars += (ra,)

        self.pr = pr
        self.ra = ra

        u = self.create_function("u", invars)
        v = self.create_function("v", invars)
        w = self.create_function("w", invars) if dim == 3 else sp.Number(0)
        p = self.create_function("p", invars)
        T = self.create_function("T", invars)

        # diff_T = T.diff(y)
        # diff_p_equation_1 = p.diff(x) - pr * u.diff(x).diff(x)
        # diff_p_equation_2 = p.diff(y) - pr * u.diff(y).diff(y) - pr * ra * T

        continuity = u.diff(x) + v.diff(y) + w.diff(z)

        momentum_x = (
                u * u.diff(x)
                + v * u.diff(y)
                + w * u.diff(z)
                - (
                        (pr * u.diff(x)).diff(x)
                        + (pr * u.diff(y)).diff(y)
                        + (pr * u.diff(z)).diff(z)
                )
                + p.diff(x)
        )

        momentum_y = (
                u * v.diff(x)
                + v * v.diff(y)
                + w * v.diff(z)
                - (
                        (pr * v.diff(x)).diff(x)
                        + (pr * v.diff(y)).diff(y)
                        + (pr * v.diff(z)).diff(z)
                )
                + p.diff(y) - pr * ra * T
        )

        energy = (
                u * T.diff(x)
                + v * T.diff(y)
                + w * T.diff(z)
                - (
                        T.diff(x).diff(x)
                        + T.diff(y).diff(y)
                        + T.diff(z).diff(z)
                )
        )

        self.add_equation("continuity", continuity)
        self.add_equation("momentum_x", momentum_x)
        self.add_equation("momentum_y", momentum_y)
        self.add_equation("energy", energy)
        # self.add_equation("diff_T", diff_T)
        # self.add_equation("diff_p_equation_1", diff_p_equation_1)
        # self.add_equation("diff_p_equation_2", diff_p_equation_2)

        self._apply_detach()


def train(cfg: DictConfig):
    # set random seed for reproducibility
    ppsci.utils.misc.set_random_seed(cfg.seed)
    # initialize logger
    logger.init_logger("ppsci", osp.join(cfg.output_dir, "train.log"), "info")

    # set model
    model = ppsci.arch.MLP(**cfg.MODEL)

    # set equation
    equation = {"CNC2D": CNC2D(cfg.pr, cfg.ra, cfg.DIMENSION)}

    # set timestamps(including initial t0)
    # timestamps = np.linspace(0.0, 1.5, cfg.NTIME_ALL, endpoint=True)
    # set time-geometry
    geom = {
        "rect": ppsci.geometry.Rectangle((0, 0), (1, 1))
    }

    # set dataloader config
    train_dataloader_cfg = {
        "dataset": "IterableNamedArrayDataset",
        "iters_per_epoch": cfg.TRAIN.iters_per_epoch,
    }

    # pde/bc constraint use t1~tn, initial constraint use t0
    # NPOINT_PDE, NTIME_PDE = 99 ** 2, cfg.NTIME_ALL - 1
    # NPOINT_TOP, NTIME_TOP = 101, cfg.NTIME_ALL - 1
    # NPOINT_DOWN, NTIME_DOWN = 101, cfg.NTIME_ALL - 1
    # NPOINT_LEFT, NTIME_LEFT = 99, cfg.NTIME_ALL - 1
    # NPOINT_RIGHT, NTIME_RIGHT = 99, cfg.NTIME_ALL - 1
    # NPOINT_IC, NTIME_IC = 99 ** 2, 1

    NPOINT_PDE = 99 ** 2
    NPOINT_TOP = 101
    NPOINT_BOTTOM = 101
    NPOINT_LEFT = 99
    NPOINT_RIGHT = 99

    # set constraint
    pde = ppsci.constraint.InteriorConstraint(
        equation["CNC2D"].equations,
        {"continuity": 0, "momentum_x": 0, "momentum_y": 0, "energy": 0},
        geom["rect"],
        {**train_dataloader_cfg, "batch_size": NPOINT_PDE},
        ppsci.loss.MSELoss("sum"),
        evenly=True,
        weight_dict=cfg.TRAIN.weight.pde,  # (1)
        name="EQ",
    )

    bc_top = ppsci.constraint.BoundaryConstraint(
        {
            "u": lambda out: out["u"],
            "v": lambda out: out["v"],
            "diff_T": lambda out: jacobian(out["T"], out["y"])[1],
            "diff_p_equation_2": lambda out: jacobian(out["p"],out["y"]) -
                                             cfg.pr * hessian(out["u"], out["y"]) -
                                             cfg.pr * cfg.ra * out["T"]
        },
        {"u": 0, "v": 0, "diff_T": 0, "diff_p_equation_2": 0},
        geom["rect"],
        {**train_dataloader_cfg, "batch_size": NPOINT_TOP},
        ppsci.loss.MSELoss("sum"),
        criteria=lambda x, y: np.isclose(y, 1),
        name="BC_top",
    )

    # TODO: 以此类推
    bc_down = ppsci.constraint.BoundaryConstraint(
        {
            "u": lambda out: out["u"],
            "v": lambda out: out["v"],
            "diff_T": lambda out: jacobian(out["T"], out["y"]),
            "diff_p_equation_2": lambda out: jacobian(out["p"],out["y"]) -
                                             cfg.pr * hessian(out["u"], out["y"]) -
                                             cfg.pr * cfg.ra * out["T"]
        },
        {"u": 0, "v": 0, "diff_T": 0, "diff_p_equation_2": 0},
        geom["rect"],
        {**train_dataloader_cfg, "batch_size": NPOINT_BOTTOM},
        ppsci.loss.MSELoss("sum"),
        criteria=lambda x, y: np.isclose(y, 0),
        name="BC_down",
    )

    bc_left = ppsci.constraint.BoundaryConstraint(
        {
            "u": lambda out: out["u"],
            "v": lambda out: out["v"],
            "T": lambda out: out["T"],
            "diff_p_equation_1": lambda out: jacobian(out["p"], out["x"]) -
                                        cfg.pr * hessian(out["u"], out["x"])
         },
        {"u": 0, "v": 0, "T": 1, "diff_p_equation_1": 0},
        geom["rect"],
        {**train_dataloader_cfg, "batch_size": NPOINT_LEFT},
        ppsci.loss.MSELoss("sum"),
        criteria=lambda x, y: np.isclose(x, 0),
        name="BC_left",
    )
    bc_right = ppsci.constraint.BoundaryConstraint(
        {
            "u": lambda out: out["u"],
            "v": lambda out: out["v"],
            "T": lambda out: out["T"],
            "diff_p_equation_1": lambda out: jacobian(out["p"], out["x"]) -
                                        cfg.pr * hessian(out["u"], out["x"])
         },
        {"u": 0, "v": 0, "T": 0, "diff_p_equation_1": 0},
        geom["rect"],
        {**train_dataloader_cfg, "batch_size": NPOINT_RIGHT},
        ppsci.loss.MSELoss("sum"),
        criteria=lambda x, y: np.isclose(x, 1),
        name="BC_right",
    )
    # ic = ppsci.constraint.InitialConstraint(
    #     {"u": lambda out: out["u"], "v": lambda out: out["v"]},
    #     {"u": 0, "v": 0},
    #     geom["rect"],
    #     {**train_dataloader_cfg, "batch_size": NPOINT_IC * NTIME_IC},
    #     ppsci.loss.MSELoss("sum"),
    #     evenly=True,
    #     name="IC",
    # )
    # wrap constraints together
    constraint = {
        pde.name: pde,
        bc_top.name: bc_top,
        bc_down.name: bc_down,
        bc_left.name: bc_left,
        bc_right.name: bc_right,
        # ic.name: ic,
    }

    # set training hyper-parameters
    lr_scheduler = ppsci.optimizer.lr_scheduler.Cosine(
        **cfg.TRAIN.lr_scheduler,
        warmup_epoch=int(0.05 * cfg.TRAIN.epochs),
    )()

    # set optimizer
    optimizer = ppsci.optimizer.Adam(lr_scheduler)(model)

    # set validator
    NPOINT_EVAL = NPOINT_PDE
    residual_validator = ppsci.validate.GeometryValidator(
        equation["CNC2D"].equations,
        {"continuity": 0, "momentum_x": 0, "momentum_y": 0, "energy": 0},
        geom["rect"],
        {
            "dataset": "NamedArrayDataset",
            "total_size": NPOINT_EVAL,
            "batch_size": cfg.EVAL.batch_size.residual_validator,
            "sampler": {"name": "BatchSampler"},
        },
        ppsci.loss.MSELoss("sum"),
        evenly=True,
        metric={"MSE": ppsci.metric.MSE()},
        name="Residual",
    )
    validator = {residual_validator.name: residual_validator}

    # set visualizer(optional)
    NPOINT_BC = NPOINT_TOP + NPOINT_BOTTOM + NPOINT_LEFT + NPOINT_RIGHT
    vis_points = geom["rect"].sample_interior(NPOINT_PDE + NPOINT_BC, evenly=True)
    visualizer = {
        "visualize_u_v": ppsci.visualize.VisualizerVtu(
            vis_points,
            {"u": lambda d: d["u"], "v": lambda d: d["v"], "p": lambda d: d["p"], "T": lambda d: d["T"]},
            prefix="result_u_v",
        )
    }

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        constraint,
        cfg.output_dir,
        optimizer,
        lr_scheduler,
        cfg.TRAIN.epochs,
        cfg.TRAIN.iters_per_epoch,
        eval_during_train=cfg.TRAIN.eval_during_train,
        eval_freq=cfg.TRAIN.eval_freq,
        equation=equation,
        geom=geom,
        validator=validator,
        visualizer=visualizer,
        checkpoint_path=cfg.TRAIN.checkpoint_path,
    )
    # train model
    solver.train()
    # evaluate after finished training
    solver.eval()
    # visualize prediction after finished training
    solver.visualize()


@hydra.main(
    version_base=None, config_path="./", config_name="cnc_2d.yaml"
)
def run(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    # elif cfg.mode == "eval":
    #     evaluate(cfg)
    # elif cfg.mode == "export":
    #     export(cfg)
    # elif cfg.mode == "infer":
    #     inference(cfg)
    else:
        raise ValueError(
            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
        )


if __name__ == "__main__":
    run()
